Skip to content

feat: bulk customer usage attribution lookup performance and correctness#4686

Merged
gergely-kurucz-konghq merged 13 commits into
mainfrom
feat/bulk-customer-usage-attribution-lookup-perf
Jul 23, 2026
Merged

feat: bulk customer usage attribution lookup performance and correctness#4686
gergely-kurucz-konghq merged 13 commits into
mainfrom
feat/bulk-customer-usage-attribution-lookup-perf

Conversation

@gergely-kurucz-konghq

@gergely-kurucz-konghq gergely-kurucz-konghq commented Jul 10, 2026

Copy link
Copy Markdown
Contributor

What

Optimizes and unifies the customer usage-attribution lookups. The bulk GetCustomersByUsageAttribution still had the cross-table OR anti-pattern that #4684 fixed for the single-key path. This PR fixes the bulk perf, collapses the single-key and bulk paths onto one predicate and one precedence mechanism, and adds layer-separated benchmarks.

Why

  • Perf: same seq-scan pathology as fix: optimize customer usage attribution lookup #4684 — Postgres couldn't index either branch of the OR. Callers can pass an uncapped key set, so this is real production exposure, not just a consistency fix.
  • Correctness: the bulk lookup had no key-over-subject precedence, unlike the single-key path. Callers resolved collisions with a single-pass map[key]customer that depended on DB row order — provably unfixable by reordering when two customers collide on the same key from opposite sides (one's own key vs. the other's subject key).
  • Consistency: single-key and bulk carried two precedence implementations (SQL-side ORDER BY lookup_priority LIMIT 1 vs. an in-memory pass) and two nearly-identical predicates — a drift hazard.

How

  • Rewrote the bulk query as two independently-indexable UNION ALL branches (same shape as fix: optimize customer usage attribution lookup #4684), returning the full candidate set.
  • Moved key resolution into customer.Service: GetCustomersByUsageAttribution now returns map[string]Customer with precedence resolved once, centrally — deleting the duplicated (and buggy) mapping logic from callers (governance).
  • Unified single-key onto the same path: deleted the single-key adapter method and its SQL-precedence helper; the single-key service method now calls the bulk adapter with a one-element key set and resolves via the shared resolveCustomersByKeyWithPrecedence. The single-key public signature and error contract are unchanged (validation / not-found / never-conflict), so no consumer changes.
  • Applied the deleted_at grace window (IS NULL OR > at) uniformly across both predicate branches, matching the WithSubjects convention. The created_at active-time guard is intentionally omitted — created_at is server-assigned and immutable, so a future-dated subject can't occur through the service; guarding only one path would introduce a batch-composition-dependent bug for a state that can't happen.
  • Dropped the ambiguous-collision logging (collisions are structurally rare and resolved deterministically; the DB can be queried directly if ever needed) and the now-unused service logger.
  • Split the benchmarks by layer — adapter (candidate query only) vs. service (full path incl. precedence + hydration) — with shared 100k-row SQL seed fixtures in openmeter/customer/testutils/uabench.
  • Breaking signature change on customer.Service.GetCustomersByUsageAttribution (now returns a map).

Benchmarks

Adapter — candidate query only, cross-table OR vs UNION ALL (100k customers, median of 3)

Route Cross-table OR UNION ALL Improvement
single-key (by customer key) 8.19 ms/op 0.40 ms/op ~20x
single-key (by subject key) 11.45 ms/op 0.46 ms/op ~25x
bulk (100 keys, mixed) 229 ms/op 0.98 ms/op ~220x

Service — full lookup path (query + subject hydration + key-over-subject precedence, 100k customers, median of 3)

Path Latency allocs/op
single-key (by customer key) 0.78 ms/op 910
single-key (by subject key) 0.78 ms/op 911
bulk (100 keys) 2.11 ms/op 6,965
POSTGRES_HOST=127.0.0.1 go test -tags=dynamic \
  ./openmeter/customer/adapter ./openmeter/customer/service \
  -run '^$' -bench 'UsageAttribution' -benchtime=10x -count=3 -benchmem

Governance endpoint — namespace-scale (fixed 100 customers × 1 feature)

decoys pre-fix (OR) post-fix (UNION ALL)
0 38.0 ms 48.4 ms
10,000 42.0 ms 50.9 ms
100,000 159.7 ms 53.8 ms

Pre-fix: 4.2x jump (38 ms → 160 ms) as namespace grows. Post-fix: 1.11x (flat, within noise).

cd e2e && make bench-governance-namespace

Governance endpoint — customers × features diagonal (baseline, no regression)

customers = features latency allocs
1 1.2 ms 117
10 9.8 ms 467
50 150 ms 6,083
100 547 ms 22,240
cd e2e && make bench-governance

Summary by CodeRabbit

Summary by CodeRabbit

  • Improvements

    • Updated usage-attribution customer lookup to return a key → customer mapping.
    • Enforced key-over-subject precedence; unmatched keys return nil entries.
    • Added stricter validation for empty key inputs.
  • Performance

    • Added benchmarks for governance namespace scaling.
    • Added benchmarks comparing bulk vs alternative customer-resolution strategies across different key sizes.
  • Developer Experience

    • Added a command/Make target to run the governance namespace benchmark.
    • Improved E2E Postgres setup for local and CI execution.

Greptile Summary

This PR improves customer usage-attribution lookup performance and centralizes resolution behavior. The main changes are:

  • Replaces the bulk cross-table lookup with indexable UNION ALL candidate queries.
  • Routes single-key and bulk lookups through one service-level precedence resolver.
  • Changes bulk service resolution to return a key-to-customer map with nil entries for misses.
  • Updates governance to consume the resolved map directly.
  • Adds adapter, service, and e2e benchmarks for the lookup path.

Confidence Score: 5/5

This looks safe to merge.

  • No blocking issues found in the changed code.

Important Files Changed

Filename Overview
openmeter/customer/adapter/customer.go Reworks the usage-attribution predicate into shared UNION ALL candidate branches.
openmeter/customer/service/customer.go Adds shared key-over-subject precedence resolution for single-key and bulk lookup paths.
openmeter/governance/service/service.go Uses the service-resolved key-to-customer map instead of rebuilding local lookup state.
openmeter/customer/service.go Updates the bulk customer service method to return resolved customer pointers by key.
e2e/governance_bench_namespace_test.go Adds a namespace-scaling benchmark for governance customer resolution.

Reviews (6): Last reviewed commit: "refactor(customer): return map[string]*C..." | Re-trigger Greptile

Context used (3)

  • Context used - CLAUDE.md (source)
  • Context used - AGENTS.md (source)
  • Context used - api/spec/AGENTS.md (source)

@gergely-kurucz-konghq
gergely-kurucz-konghq requested a review from a team as a code owner July 10, 2026 17:23
@gergely-kurucz-konghq gergely-kurucz-konghq added release-note/misc Miscellaneous changes kind/refactor Code refactor, cleanup or minor improvement labels Jul 10, 2026
@coderabbitai

coderabbitai Bot commented Jul 10, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

📝 Walkthrough

Walkthrough

Usage-attribution resolution now returns key-to-customer maps with key-over-subject precedence. Governance consumes the resolved map directly, while adapter, service, and E2E benchmarks measure lookup and namespace-scaling behavior.

Changes

Usage attribution resolution

Layer / File(s) Summary
Bulk customer resolution
openmeter/customer/adapter/customer.go, openmeter/customer/service.go, openmeter/customer/service/customer.go, openmeter/customer/service/customer_test.go, openmeter/customer/service/service_test.go
Bulk matching uses customer and subject keys, and the service resolves requested keys with customer-key precedence while returning nil for unmatched keys.
Governance resolution integration
openmeter/governance/service/service.go, openmeter/server/server_test.go
Governance consumes the resolved customer map directly, and the no-op customer service matches the updated interface.
Lookup benchmark coverage
openmeter/customer/adapter/customer_usage_attribution_benchmark_test.go, openmeter/customer/service/customer_usage_attribution_benchmark_test.go
Benchmarks cover single-key and bulk lookup paths, including cross-table and UNION ALL strategies.
E2E namespace benchmark
e2e/governance_bench_namespace_test.go, e2e/helpers.go, e2e/ledger_accounts_test.go, e2e/Makefile
A benchmark command seeds namespace decoys incrementally and measures governance query latency using shared Postgres setup.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant GovernanceService
  participant CustomerService
  participant CustomerAdapter
  participant Postgres
  GovernanceService->>CustomerService: Resolve usage-attribution keys
  CustomerService->>CustomerAdapter: Fetch matching customers
  CustomerAdapter->>Postgres: Query customer and subject-key candidates
  Postgres-->>CustomerAdapter: Return candidate customer data
  CustomerAdapter-->>CustomerService: Return candidates
  CustomerService->>CustomerService: Apply key-over-subject precedence
  CustomerService-->>GovernanceService: Return key-to-customer map
Loading

Possibly related PRs

Suggested labels: area/customers

Suggested reviewers: turip, chrisgacsal

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 44.12% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly matches the main change: bulk customer usage-attribution lookup improvements in both performance and correctness.
✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch feat/bulk-customer-usage-attribution-lookup-perf

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

Comment thread openmeter/customer/adapter/customer.go Outdated
Comment thread openmeter/customer/adapter/customer.go Outdated

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@e2e/governance_bench_namespace_test.go`:
- Line 57: Validate the custom value returned by
governanceBenchNamespaceDecoyCount before constructing decoyCounts, requiring it
to be at least 10,000 (and fail clearly if not). Update the setup around the
decoy-count benchmark cases so decoyCounts remains sorted and the incremental
seeding logic reports accurate metrics.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 8a2301f5-8bfa-48bc-af01-381cec6cb90a

📥 Commits

Reviewing files that changed from the base of the PR and between 12ab7b0 and c38b29b.

📒 Files selected for processing (27)
  • app/common/customer.go
  • e2e/Makefile
  • e2e/governance_bench_namespace_test.go
  • e2e/helpers.go
  • e2e/ledger_accounts_test.go
  • openmeter/customer/adapter/customer.go
  • openmeter/customer/adapter/customer_usage_attribution_benchmark_test.go
  • openmeter/customer/service.go
  • openmeter/customer/service/customer.go
  • openmeter/customer/service/customer_test.go
  • openmeter/customer/service/hooks/subjectcustomer_test.go
  • openmeter/customer/service/service.go
  • openmeter/customer/testutils/env.go
  • openmeter/entitlement/metered/lateevents_test.go
  • openmeter/entitlement/metered/utils_test.go
  • openmeter/entitlement/service/utils_test.go
  • openmeter/governance/service/service.go
  • openmeter/governance/service/service_test.go
  • openmeter/server/server_test.go
  • openmeter/subject/service/hooks/customersubject_test.go
  • openmeter/subject/testutils/env.go
  • openmeter/subscription/testutils/customer.go
  • test/app/stripe/testenv.go
  • test/app/testenv.go
  • test/billing/suite.go
  • test/customer/testenv.go
  • test/entitlement/regression/framework_test.go
💤 Files with no reviewable changes (1)
  • e2e/ledger_accounts_test.go

Comment thread e2e/governance_bench_namespace_test.go
@greptile-apps

greptile-apps Bot commented Jul 13, 2026

Copy link
Copy Markdown
Contributor

Want your agent to iterate on Greptile's feedback? Try greploops.

Comment thread openmeter/customer/service/customer.go Outdated
GetCustomersByUsageAttribution had the same cross-table OR anti-pattern
PR #4684 fixed for the single-key lookup, causing a Postgres seq scan on
large namespaces. Split into two independently-indexable UNION ALL
branches, same as the single-key fix, without lookup_priority (bulk still
returns every match). Benchmark: 222ms -> 1ms (220x) at 100 keys / 100k
customers.
…ally

GetCustomersByUsageAttribution returned an unordered []Customer, so
callers (governance + a private-repo consumer) each built their own
map[key]Customer with first-match-wins semantics that depended on DB row
order — unfixable by reordering when two customers collide on the same
key from opposite sides (one's own key vs. the other's subject key).

Service now returns map[string]Customer with key-over-subject precedence
resolved once, centrally, mirroring the single-key lookup's UNION ALL
priority. Each ambiguous collision is logged with the key and both
customer IDs for operator follow-up, replacing the duplicated (and now
provably buggy) mapping loops in every consumer.
Isolates namespace size from the customers/features axes
BenchmarkGovernanceQuery already covers: fixes query load at 100
customers x 1 feature, varies decoy customers/subjects seeded directly
via SQL (HTTP creation would dominate setup time at 100k+). Confirmed
against the pre-fix bulk lookup: latency scales with namespace size
(38ms -> 160ms at 0 -> 100k decoys); post-fix it stays flat (48ms ->
54ms), matching the adapter-level benchmark's win.

Moves initE2EPostgresPool into helpers.go so both this and the
existing ledger e2e check can share it.
… size

decoyCounts hardcoded 10_000 as the middle tier; GOV_BENCH_NAMESPACE_DECOYS
below that seeded up to 10_000 first, then skipped seeding for the smaller
tier (seeded > target) while still reporting the smaller, now-incorrect
count as the benchmark label. Sort + dedup keeps the sequence monotonic
regardless of the env override.
…ibution lookup

Bulk subject-key matching added a CreatedAtLTE guard the single-key path
never had, diverging from its own eager-loaded subject list and forcing a
second query to reconcile them. created_at is server-assigned and immutable,
so a future-created subject can't occur in production. Remove the guard to
restore parity with the single-key predicate.
Trim redundant clauses in two comments (governance namespace-scale
benchmark doc, bulk usage-attribution created_at-guard note) that
restated the same point twice or over-qualified it. No behavior change.
…bution lookups

Give the subject-branch owning-customer check the same Or(IsNull, GT(at))
grace window as the key branch in both GetCustomerByUsageAttribution and
GetCustomersByUsageAttribution, so a customer resolves identically via its
own key or a subject key. Behavior-equivalent in production (deletion always
stamps clock.Now().UTC()); adds CustomerDeletedInFutureIncluded tests.
…resolve precedence in service

Delete the single-key adapter GetCustomerByUsageAttribution and its SQL-precedence
helper (lookup_priority / ORDER BY / LIMIT). Both single-key and bulk lookups now use
the one customersMatchUsageAttributionKeys predicate; the single-key service method
calls the bulk adapter with a one-element key set and applies key-over-subject
precedence in Go via resolveCustomersByKey. The service signature and error contract
are unchanged (validation / not-found-same-message / never-conflict), so no OSS or
downstream consumer changes.

Drop the ambiguous-match logging from both single and bulk paths: key-over-subject
collisions are structurally rare and resolved deterministically, so resolveCustomersByKey
now returns only the resolved map. A benchmark of the single-key path (SQL-side vs
Go-side precedence) confirms it is perf-neutral, with slightly fewer allocations.
…ution core

Route the single-key and bulk GetCustomer(s)ByUsageAttribution service methods
through one resolveCustomersByUsageAttribution helper (adapter fetch + precedence)
so their core resolution can't drift. Rename resolveCustomersByKey to
resolveCustomersByKeyWithPrecedence to name its key-over-subject behavior.

Add Test_GetCustomersByUsageAttribution covering key/subject resolution,
key-over-subject precedence, unmatched-key-absent (no error), and empty-key
validation, mirroring Test_GetCustomerByUsageAttribution. The fully symmetric
{K1,K2} cross-collision is left to the unit and adapter tests: CreateCustomer's
overlap guard rejects a customer key that overlaps another customer's subject,
so that state can't be constructed through the service.
The adapter benchmark inlined a copy of the service's key-over-subject
precedence, mixing layers and risking drift from resolveCustomersByKeyWithPrecedence.
Split into layer-honest benchmarks:

- adapter package: BenchmarkCustomer(s)UsageAttribution*Query measure only the
  candidate query (cross_table_or vs union_all), no precedence.
- service package: BenchmarkCustomer(s)UsageAttribution*Lookup drive the real
  service.Get(s)ByUsageAttribution end-to-end, exercising the real precedence.

Shared seed/count/bulk-key fixtures move to a new leaf package
customer/testutils/uabench (imports only the ent client, so both benchmark
packages use it without an import cycle). The inline precedence copy is deleted;
the seed SQL is unchanged.

Run the timed loop of the bulk service benchmark inside b.Run so it gets a
fresh, running timer: the previous flat form called b.StopTimer() during setup
and b.ResetTimer() (which does not restart a stopped timer), so the loop was
measured with the timer off and reported no ns/op.

Test-only, no production change.
Rebase onto main changed the harness: NewTestEnv migrates internally and
InitPostgresDB takes PostgresDBStateEntMigrated. Drop the leftover
DBSchemaMigrate call and stale comment, and the now-redundant Schema.Create
in the adapter and service benchmarks.
@gergely-kurucz-konghq
gergely-kurucz-konghq force-pushed the feat/bulk-customer-usage-attribution-lookup-perf branch from e8ddca3 to 84e266c Compare July 17, 2026 15:15
The service Logger was added on this branch only for the ambiguous-match
logging, which was removed; s.logger had no remaining uses. Remove the field
from customerservice.Config and the Logger argument from all call sites.
Comment thread openmeter/customer/adapter/customer.go
@gergely-kurucz-konghq

Copy link
Copy Markdown
Contributor Author

@coderabbitai resume

@coderabbitai

coderabbitai Bot commented Jul 17, 2026

Copy link
Copy Markdown
Contributor
✅ Action performed

Reviews resumed.

turip
turip previously approved these changes Jul 21, 2026
…ageAttribution

Align the bulk resolver's semantics with the single-key lookup and the repo's
featureresolver.BatchResolve convention: every input key is present in the
returned map, with a nil value marking not-found (was: value map, absent =
not-found). Callers can distinguish requested-but-missing from not-requested,
and the map holds pointers instead of copied Customer structs.
@gergely-kurucz-konghq
gergely-kurucz-konghq merged commit 8cdf8ad into main Jul 23, 2026
27 checks passed
@gergely-kurucz-konghq
gergely-kurucz-konghq deleted the feat/bulk-customer-usage-attribution-lookup-perf branch July 23, 2026 09:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

kind/refactor Code refactor, cleanup or minor improvement release-note/misc Miscellaneous changes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants